import datetime
import json
import os
import pickle

import fire
import torch
import torch.utils.data as torchdata
import numpy as np

from codesecurity.interface.api import ObfuscatorMode, obfuscate_js,jstap_train,JStapMeta
from codesecurity.tasks.malicious_code_detect.mcd import (build_ast_caches,build_source_image,
build_dataset,training_mcd_model,eval_mcd_model,training_eval_mcd_model,
build_ast_image,build_ast_image_without_labels,build_dataset_without_labels,
build_ast_caches_without_labels,build_source_image_without_labels,execute_mcd_model,build_ast_caches_without_labels_with_history)
from codesecurity.tasks.malicious_code_detect.meta import SuperParameter_MCD,ModelCachesMeta


def obfuscate(in_dir, out_dir, mode:str):
    obfuscate_js(in_dir, out_dir, ObfuscatorMode.instance(mode))


def train_jstap(dataset_dir,rebuild=False):
    meta=JStapMeta.auto(dataset_dir)
    jstap_train(meta,rebuild)
# python mcd_main.py obfuscate --in_dir third_party/JStap/samples/Bad-validate --out_dir data/js/obfucscate/light/jstap_validate/bad --mode light

def build_ast(dataset_dir,parallel=False,without_label=False,caches_file=None):
    if without_label:
        if caches_file is None:
            dir_name=os.path.basename(dataset_dir)
            caches_file=os.path.join('model_caches/data_caches',f'{dir_name}_ast_caches.pt')
        build_ast_caches_without_labels(dataset_dir,caches_file,parallel)
    else:
        build_ast_caches(dataset_dir,'good','bad',parallel)

def build_ast_robustness(dataset_dir,caches_file=None):
    if caches_file is None:
            dir_name=os.path.basename(dataset_dir)
            caches_file=os.path.join('model_caches/data_caches',f'{dir_name}_ast_caches.pt')
    build_ast_caches_without_labels_with_history(dataset_dir,caches_file)

def build_image_ast(dataset_dir,parallel=False,without_label=False,caches_file=None):
    sp=SuperParameter_MCD.default()
    
    if without_label:
        ast_caches=dataset_dir
        if caches_file is None:
            ast_caches_name=os.path.basename(ast_caches)
            caches_name=ast_caches_name.replace('_ast_caches','_ast_image_caches')
            caches_file=os.path.join('model_caches/data_caches',f'{caches_name}')
        build_ast_image_without_labels(ast_caches,sp,caches_file,parallel)
    
    else:
        build_ast_image(dataset_dir,sp,'good','bad',parallel)


def build_image(dataset_dir,parallel=False,without_label=False,caches_file=None):
    sp=SuperParameter_MCD.default()
    
    if without_label:
        ast_caches=dataset_dir
        if caches_file is None:
            ast_caches_name=os.path.basename(ast_caches)
            caches_name=ast_caches_name.replace('_ast_caches','_image_caches')
            caches_file=os.path.join('model_caches/data_caches',f'{caches_name}')
        build_source_image_without_labels(ast_caches,sp,caches_file,parallel)
    else:
        build_source_image(dataset_dir,sp,'good','bad',parallel)

def build_model(dataset_dir,eval_dataset=None,model_path=None,include_ast=False,include_dfg=True):
    
    if model_path is None:
        model_path=ModelCachesMeta.default('model_cahces/default_mcd_model').model_file
    
    include_features=[]
    if include_dfg:
        include_features.append('dfg')
    if include_ast:
        include_features.append('ast')
    
    
    sp=SuperParameter_MCD.default()
    dataset=build_dataset(dataset_dir,'good','bad',include_features=include_features)
    
    
    if eval_dataset is None:
        training_mcd_model(dataset,sp,model_path)
    else:
        training_eval_mcd_model(dataset,eval_dataset,sp,model_path)
    

def eval_model(dataset_dir,model_path=None,include_ast=False,include_dfg=True):
    if model_path is None:
        model_path=ModelCachesMeta.default('model_cahces/default_mcd_model').model_file
    
    include_features=[]
    if include_dfg:
        include_features.append('dfg')
    if include_ast:
        include_features.append('ast')
    
    
    sp=SuperParameter_MCD.default()
    dataset=build_dataset(dataset_dir,'good','bad',include_features=include_features)
    eval_mcd_model(dataset,sp,model_path)

def execute_model(model_path,*caches_pairs):

    status=[0,0]
    
    caches_pairs=[caches_pairs]
    
    sp=SuperParameter_MCD.default()
    dataset=build_dataset_without_labels(caches_pairs,['origin','dfg','ast'],True)
    dataset.features.label_with_origin()
    result=execute_mcd_model(model_path,sp,dataset)
    for i,e in enumerate(result):
        x,y,y_hat=e
        if y_hat[0]>y_hat[1]:
            status[0]+=1
        else:
            status[1]+=1
            print(y)
    print(f'good:{status[0]} bad:{status[1]}')
    
def m(magic_word):
    target_dirs=['data/js/jstap_train','data/js/obfucscate/light/jstap_train','data/js/obfucscate/medium/jstap_train']
    parallel=False
    magic_word=str(magic_word)
    
    
    target_dir=target_dirs[0]
    
    if 'e' in magic_word:
        target_dir=target_dirs[2]
        
    if 't' in magic_word:
        target_dir=target_dirs[1]
    
    if 'p' in magic_word:
        parallel=True
    
    include_ast=False
    include_dfg=True
    
    if 'a' in magic_word:
        include_ast=True
        include_dfg=False
        
    if 'b' in magic_word:
        include_ast=True
        include_dfg=True
    
    if '1' in magic_word:
        build_ast(target_dir,parallel=parallel)
    if '2' in magic_word:
        build_image(target_dir,parallel=parallel)
    if '3' in magic_word:
        build_model(target_dir,include_ast=include_ast,include_dfg=include_dfg)
    if '4' in magic_word:
        eval_model(target_dir,include_ast=include_ast,include_dfg=include_dfg)
# python mcd_main.py build_ast data/js/obfucscate/medium/jstap_train
# python mcd_main.py build_image data/js/obfucscate/medium/jstap_train
# python mcd_main.py build_model data/js/jstap_train
# python mcd_main.py eval_model data/js/obfucscate/medium/jstap_train
# python mcd_main.py build_image_ast data/js/obfucscate/medium/jstap_train
    
if __name__=="__main__":
    fire.Fire()
